import cv2
import os
from tqdm import tqdm
import mediapipe as mp

# 参数
OUTPUT_NAME_PADDING = 3  # 输出文件名帧号补零
FRAME_STEP = 1           # 每隔多少帧检测一次

# 初始化 mediapipe Hands
mp_hands = mp.solutions.hands
hands_detector = mp_hands.Hands(
    static_image_mode=True,  # 单帧模式
    max_num_hands=2,
    min_detection_confidence=0.5
)

def hands_overlap(frame):
    """
    用 mediapipe 检测是否有手
    """
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = hands_detector.process(rgb)
    return results.multi_hand_landmarks is not None

def process_video(video_path, output_root):
    basename = os.path.splitext(os.path.basename(video_path))[0]
    out_dir = os.path.join(output_root, basename + "_folder")
    os.makedirs(out_dir, exist_ok=True)

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"⚠️ 无法打开视频: {video_path}")
        return

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    saved_frames = []

    frame_idx = 0
    pbar = tqdm(total=frame_count, desc=basename)
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_idx % FRAME_STEP == 0:
            if not hands_overlap(frame):
                saved_frames.append((frame_idx, frame.copy()))

        frame_idx += 1
        pbar.update(1)
    pbar.close()
    cap.release()

    # -------- 选择 4 张帧 --------
    selected = []
    if len(saved_frames) >= 4:
        selected = [
            saved_frames[0],                          # 最早无手（0 积木）
            saved_frames[len(saved_frames)//3],       # 约 1/3 位置（2 积木）
            saved_frames[2*len(saved_frames)//3],     # 约 2/3 位置（4 积木）
            saved_frames[-1]                          # 最后无手（6 积木）
        ]
    else:
        # 如果不足 4 张，尽量补齐，最后一张重复
        selected = saved_frames
        while len(selected) < 4 and len(saved_frames) > 0:
            selected.append(saved_frames[-1])

    # -------- 保存 --------
    for i, (frame_number, frame) in enumerate(selected):
        fname = str(frame_number).zfill(OUTPUT_NAME_PADDING) + ".png"
        out_path = os.path.join(out_dir, fname)
        cv2.imwrite(out_path, frame)
        print(f"保存 {i+1}/4: {out_path}")

    print(f"完成：{basename}, 共保存 {len(selected)}/4 张")

if __name__ == "__main__":
    video_folder = "/home/mk/视频/newhand/output"  # H.265 视频文件夹
    output_root = "outputd"
    os.makedirs(output_root, exist_ok=True)

    for fname in os.listdir(video_folder):
        if fname.lower().endswith(".mp4"):
            process_video(os.path.join(video_folder, fname), output_root)

    hands_detector.close()
